#include "NCNNModel.hpp"
#include <iostream>
#include <chrono>
#include <opencv2/opencv.hpp>

// using clock_t = std::chrono::steady_clock;
using duration_t = std::chrono::duration<double, std::milli>;
struct TimingStats;
TimingStats analyze_timings(const std::vector<double>& timings);
void print_performance_report(
    const std::vector<double>& timings_preprocess,
    const std::vector<double>& timings_infer,
    const std::vector<double>& timings_total);
void print_time_distribution_aligned(const std::vector<double>& timings, std::string desc = "", double bin_size = 0.0);

int main() {
    // 模型路径
    std::string model_param = "../../yolov8n_ncnn_model/model.ncnn.param";
    std::string model_bin   = "../../yolov8n_ncnn_model/model.ncnn.bin";
    cv::VideoCapture cap(0);


    // 设置输入输出名称（YOLOv8/YOLO11 常见格式）
    std::string input_name = "in0";
    std::vector<std::string> output_names = {"out0"};  // 多头输出可加多个

    // 创建模型实例
    NCNNModel model(model_param, model_bin, 640, input_name, output_names, true);

    // 读取图像
    cv::Mat image = cv::Mat::zeros(cv::Size(640, 640), CV_8UC3);
    if (image.empty()) {
        std::cerr << "Failed to load image!" << std::endl;
        return -1;
    }
    ncnn::Mat input = model.preprocess(image);

    // 预热阶段
    int warmup_iterations = 500;
    std::cout << "Warming up for " << warmup_iterations << " iterations..." << std::endl;
    {
        for (int i = 0; i < warmup_iterations; ++i) {
            ncnn::Mat input = model.preprocess(image);
            model.infer(input);
        }
    }

    // 正式推理阶段
    int test_iterations = 100;
    std::vector<double> timings_preprocess;
    std::vector<double> timings_infer;
    std::vector<double> timings_total;
    std::cout << "Running inference for " << test_iterations << " iterations..." << std::endl;
    {

        for (int i = 0; i < test_iterations; ++i) {
            cap >> image;
            auto t_1 = std::chrono::steady_clock::now();
            ncnn::Mat input = model.preprocess(image);
            auto t_2 = std::chrono::steady_clock::now();
            std::vector<ncnn::Mat> outputs = model.infer(input);

            auto t_3 = std::chrono::steady_clock::now();
            double preprocess_ms = std::chrono::duration_cast<duration_t>(t_2 - t_1).count();
            double infer_ms = std::chrono::duration_cast<duration_t>(t_3 - t_2).count();
            double total_ms = std::chrono::duration_cast<duration_t>(t_3 - t_1).count();
            timings_preprocess.push_back(preprocess_ms);
            timings_infer.push_back(infer_ms);
            timings_total.push_back(total_ms);
        }
    }

    // 统计结果
    print_performance_report(timings_preprocess, timings_infer, timings_total);
    // 时间分布直方图
    print_time_distribution_aligned(timings_total, "total");
    print_time_distribution_aligned(timings_infer, "infer");

    return 0;
}


struct TimingStats {
    double avg_time;
    double std_dev;
    double min_time;
    double max_time;
    double fps;  // 只有每帧有意义
};

TimingStats analyze_timings(const std::vector<double>& timings) {
    if (timings.empty()) {
        return {0, 0, 0, 0, 0};
    }

    double total_time = 0.0;
    double min_time = timings[0];
    double max_time = timings[0];

    for (double t : timings) {
        total_time += t;
        if (t < min_time) min_time = t;
        if (t > max_time) max_time = t;
    }

    double avg_time = total_time / timings.size();

    double variance = 0.0;
    for (double t : timings) {
        variance += (t - avg_time) * (t - avg_time);
    }
    variance /= timings.size();
    double std_dev = std::sqrt(variance);

    double fps = (avg_time > 0) ? (1000.0 / avg_time) : 0.0;

    return {avg_time, std_dev, min_time, max_time, fps};
}

// 将数值格式化为固定宽度字符串
template <typename T>
std::string format_fixed_width(T value, int width, int precision = -1) {
    std::ostringstream oss;
    if (precision >= 0) {
        oss << std::fixed << std::setprecision(precision);
    }
    oss << value;
    std::string s = oss.str();
    if (s.length() > static_cast<size_t>(width)) return s.substr(0, width);
    return s + std::string(width - s.length(), ' ');
}

std::string to_string_with_precision(double value, int precision = 2) {
    std::ostringstream oss;
    oss << std::fixed << std::setprecision(precision) << value;
    return oss.str();
}

void print_performance_report(
    const std::vector<double>& timings_preprocess,
    const std::vector<double>& timings_infer,
    const std::vector<double>& timings_total)
{
    TimingStats stats_pre = analyze_timings(timings_preprocess);
    TimingStats stats_inf = analyze_timings(timings_infer);
    TimingStats stats_tot = analyze_timings(timings_total);

    const int label_width = 15;
    const int num_width = 6;

    std::cout << "\n================ Performance Report ================" << std::endl;
    std::cout << std::left
              << std::setw(label_width) << "Stage"
              << " | "
              << std::setw(num_width) << "Avg (ms)"
              << " | "
              << std::setw(num_width) << "StdDev"
              << " | "
              << std::setw(12) << "Min ~ Max"
              << " | "
              << std::setw(6) << "FPS"
              << std::endl;

    std::cout << "--------------------------------------------------" << std::endl;

    auto print_row = [&](const TimingStats& stats, const std::string& name) {
        std::cout
            << std::left << std::setw(label_width) << name << " | "
            << format_fixed_width(stats.avg_time, num_width, 2) << " | "
            << format_fixed_width(stats.std_dev, num_width, 2) << " | "
            << format_fixed_width(to_string_with_precision(stats.min_time, 2) + " ~ " + to_string_with_precision(stats.max_time, 2), 12) << " | "
            << format_fixed_width(stats.fps, 6, 2)
            << std::endl;
    };

    print_row(stats_pre, "Preprocessing");
    print_row(stats_inf, "Inference");
    print_row(stats_tot, "Total");

    std::cout << "==================================================" << std::endl;
}

// 自动计算 bin size（简单策略：取值域 / 10 的幂次，保留一位有效数字）
double auto_bin_size(double min, double max) {
    if (max <= min) return 1.0;
    double range = max - min;
    int order = static_cast<int>(floor(log10(range)));
    double base = pow(10, order);
    return std::ceil(range / base / 10) * base;
}

// 将时间分桶，并统计每段的数量
std::map<double, int> build_histogram(const std::vector<double>& timings, double bin_size, double& min_time, double& max_time) {
    if (timings.empty()) return {};

    min_time = *std::min_element(timings.begin(), timings.end());
    max_time = *std::max_element(timings.begin(), timings.end());

    std::map<double, int> histogram;

    for (double t : timings) {
        double key = bin_size * floor(t / bin_size);
        histogram[key]++;
    }

    return histogram;
}

// 合并低频、连续区间（合并 count <= threshold 的连续区间）
std::vector<std::pair<double, double>> merge_bins(
    const std::map<double, int>& histogram,
    double bin_size,
    int merge_threshold = 1)
{
    std::vector<std::pair<double, int>> flat_list(histogram.begin(), histogram.end());
    std::vector<std::pair<double, double>> merged_ranges;

    for (size_t i = 0; i < flat_list.size(); ) {
        double start = flat_list[i].first;
        int total_count = flat_list[i].second;
        size_t j = i + 1;

        while (j < flat_list.size() && flat_list[j].first == flat_list[j - 1].first + bin_size) {
            if (flat_list[j].second > merge_threshold) break;
            total_count += flat_list[j].second;
            ++j;
        }

        double end = flat_list[j - 1].first + bin_size;
        merged_ranges.emplace_back(start, end);
        i = j;
    }

    return merged_ranges;
}

// 打印时间分布表格（对齐+2位小数）
void print_time_distribution_aligned(const std::vector<double>& timings, std::string desc, double bin_size) {
    if (timings.empty()) {
        std::cout << "No timing data provided." << std::endl;
        return;
    }

    double total = timings.size();
    double min_time, max_time;

    if (bin_size <= 0) {
        bin_size = auto_bin_size(*std::min_element(timings.begin(), timings.end()),
                                 *std::max_element(timings.begin(), timings.end()));
    }

    auto histogram = build_histogram(timings, bin_size, min_time, max_time);
    auto merged_ranges = merge_bins(histogram, bin_size, 1); // 合并数量 <=1 的连续区间

    std::cout << "\n" << desc << " Time Distribution (ms):" << std::endl;
    std::cout << std::fixed << std::setprecision(2);

    // 表头
    std::cout << std::setw(13) << std::left << "Range"
              << " | "
              << std::setw(6) << std::right << "Count"
              << " | "
              << std::setw(12) << std::right << "Percentage (%)"
              << std::endl;

    std::cout << "---------------------------------------------" << std::endl;

    for (const auto& range : merged_ranges) {
        double start = range.first;
        double end = range.second;

        int count = 0;
        for (double t = start; t < end; t += bin_size) {
            auto it = histogram.find(t);
            if (it != histogram.end()) count += it->second;
        }

        double percent = count * 100.0 / total;

        char buffer[64];
        snprintf(buffer, sizeof(buffer), "[%.2f, %.2f)", start, end);

        std::cout << std::setw(13) << std::left << buffer
                  << " | "
                  << std::setw(6) << std::right << count
                  << " | "
                  << std::setw(12) << std::right << percent
                  << std::endl;
    }
}